package io.kubesphere.devops.core;

import java.util.HashMap;
import java.util.Map;

import javax.sql.DataSource;

import com.alibaba.druid.pool.DruidDataSource;

import org.mybatis.spring.SqlSessionFactoryBean;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.util.StringUtils;

@Configuration
public class DataSourceProxyConfig {

    @Primary
    @Bean("dynamicDataSource")
    public DataSource dynamicDataSource(MultipleDataSourceProperties multipleDataSourceProperties) {

        DynamicRoutingDataSource dynamicRoutingDataSource = new DynamicRoutingDataSource();
        Map<Object, Object> dataSourceMap = new HashMap<>(3);
        dataSourceMap.put(PartitionZone.ORGIN_MASTER.name(), createDataSource(multipleDataSourceProperties.getMaster(), null));
        //demo 创建两个分片库
        DataSourceProperty zone1 = multipleDataSourceProperties.getDataSourceProperty("zone1");
        if (zone1.getContainSchemas() != null) {
            zone1.getContainSchemas().stream().forEach(a ->
                    dataSourceMap.put(PartitionZone.ZONE1.name() + a, new DataSourceProxy(createDataSource(zone1, a), PartitionZone.ZONE1.name() + a, zone1.getPartitionIndex()))
            );
        }
        DataSourceProperty zone2 = multipleDataSourceProperties.getDataSourceProperty("zone2");
        if (zone2.getContainSchemas() != null) {
            zone2.getContainSchemas().stream().forEach(a ->
                    dataSourceMap.put(PartitionZone.ZONE2.name() + a, new DataSourceProxy(createDataSource(zone2, a), PartitionZone.ZONE2.name() + a, zone2.getPartitionIndex()))
            );
        }
        dynamicRoutingDataSource.setTargetDataSources(dataSourceMap);
        DynamicDataSourceContextHolder.getDataSourceKeys().addAll(dataSourceMap.keySet());
        return dynamicRoutingDataSource;
    }

    public DataSource createDataSource(DataSourceProperty dataSourceProperty, String changeSchema) {
        String url = dataSourceProperty.getUrl();
        if (changeSchema != null) {
            url = dataSourceProperty.getUrl().replace("merge_log", changeSchema);
        }
        DruidDataSource dataSource = new DruidDataSource();
        dataSource.setUsername(dataSourceProperty.getUsername());
        dataSource.setPassword(dataSourceProperty.getPassword());
        dataSource.setUrl(url);
        dataSource.setName(dataSourceProperty.getPoolName());
        String driverClassName = dataSourceProperty.getDriverClassName();
        if (!StringUtils.isEmpty(driverClassName)) {
            dataSource.setDriverClassName(driverClassName);
        }
        return dataSource;
    }

    @Bean
    @ConfigurationProperties(prefix = "mybatis")
    public SqlSessionFactoryBean sqlSessionFactoryBean(@Qualifier("dynamicDataSource") DataSource dataSource) {
        SqlSessionFactoryBean sqlSessionFactoryBean = new SqlSessionFactoryBean();
        sqlSessionFactoryBean.setDataSource(dataSource);
        return sqlSessionFactoryBean;
    }

}
